package Indexing;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

/**
 * The collection is split into several part and then processed batch by batch.
 * 
 * There are three result files:
 * 1) data//trectext.docID: integer docID - String docID
 * 2)
 * 
 * @author florahan
 */
public class MyIndexWriter {
	
	// I SUGGEST YOU TO WRITE VERY EFFICIENT CODE HERE. OTHERWISE, YOUR MEMORY CANNOT HOLD OUR CORPUS...
	private static final int MAX_PORCESSED_DOCS_PER_BATCH = 40000; //The maximum number of docs processed per batch
	private int countOfDocsPerBatch; // Count the number of docs processed per batch
	private int countOfBatches; // Count the number of batches we need to process 

	private String dataType; // The type of collection data: "trectext" or "trecweb"
	private Index index ; // The basic structure of the index. Including the ditionary and its corresponding postings lists
	private int intDocID; // Unique integer docId assigned to each document

	private BufferedWriter bwForDocIDFile; // BufferedWrite stream for writing dataType.docID file
	private BufferedWriter bwForIndexFile; // BufferedWrite stream for writing dataType.index file
	
	/**
	 * Constructor: Initate the IndexWriter
	 * @param type
	 * @throws IOException
	 */
	public MyIndexWriter(String type) throws IOException {
		// This constructor should initiate the FileWriter to output your index files.
		// Remember to close files if you finish writing the index.
		countOfDocsPerBatch = 0;
		countOfBatches = 0;
		dataType = type; 
		index = new Index(); 
		intDocID = 0;
		
	    FileOutputStream fos = new FileOutputStream("data//" + dataType + ".docID"); // Equals to "data//." + dataType + ".docID"
	    OutputStreamWriter osw = new OutputStreamWriter(fos);
	    bwForDocIDFile = new BufferedWriter(osw); 
	}
	
	/**
	 * Index a document.
	 * @param docno
	 * @param content
	 * @throws IOException
	 */
	public void IndexADocument(String docno, char[] content) throws IOException {
		// You are strongly suggested to build the index by installments
		// You need to assign the new non-negative integer docId to each document, which will be used in MyIndexReader
		if(docno == null || content == null) {
			throw new IllegalArgumentException("Invalid doc number or doc content.");
		}
		
		Map<String, ArrayList<Integer>> invertedIndexMap = new HashMap<>(); // Dictionary   Postings list(just docIDs)
																								      // apple        1 2 4 5 6 7 
																									  // bee          1 2 3 4 
		// 1. Get the content of each document																			     
        String strContent = String.valueOf(content); // Convert the char array type of content into String type of content
		
        // 2. Create inverted index for each term in the current document
        invertedIndexMap = createInvertedIndex(strContent);
		
		// 3. Update the current document into the whole index structure
		index.update(invertedIndexMap); // add inverted doc into index block
		
		// 4. Write the "integer docID - String docNO" pair into ".docID" file 
		bwForDocIDFile.write(intDocID + "," + docno + "\n");
		
		countOfDocsPerBatch++; // Count total docs indexed in ridxBlock
		intDocID++; // Increment the integer docID for next document
		
	    // 5. Write one batch of docs out into the disk once reach the max processed docs per batch
		if (countOfDocsPerBatch == MAX_PORCESSED_DOCS_PER_BATCH) { 
			writeABatchIntoDisk();
		}
	}
	/*
	 * intDocID   docNO
	 * 0          list-0000-00000
	 * 1          list-0000-00001
	 * ...        ...
	 */
	
	/**
	 * Create inverted index for each term in a document.
	 * @param content
	 * @return
	 */
	private Map<String, ArrayList<Integer>> createInvertedIndex(String content) {
		if(content == null) {
			throw new IllegalArgumentException("Invalid doc content.");
		}
		
	    Map<String, ArrayList<Integer>> InvertedIndexMap = new HashMap<>();
	    
	    // 1. Split the content of the current doc into tokens by whitespace
	    String[] tokens = content.split("\\s"); 
	    
	    // 2. Create inverted index for each term in a doc
	    for(int i = 0; i < tokens.length; i++) { // i: Track every token's position in the current doc
	        String curToken = tokens[i]; // curToken: Track the current token
	        if(!InvertedIndexMap.containsKey(curToken)) { // If the current token don't appear in the map before
	        	InvertedIndexMap.put(curToken, new ArrayList<Integer>(Arrays.asList(intDocID, 1, i))); // Put the "curToken - [intDocID, tf, positions]" pair into the InvertedIndexMap
	    	                                                                                                                 
	        } else { // If the current token has already existed in the InvertedIndexMap
	        	InvertedIndexMap.get(curToken).add(i); // Append the current position into corresponding postings list        
	    	    InvertedIndexMap.get(curToken).set(1, InvertedIndexMap.get(curToken).get(1) + 1); // Current token's tf +1
	      }
	    }
	    
	    return InvertedIndexMap;
	}
	/*
	 * [intDocID, 1, i]
     * - intDocID: 0-based incremental count of doc
     * - 1: tf
     * - i: The position of the current token in the current doc(original doc)
     * Thus, the structure of InvertedIndexMap(same as indexMap):
     * Key       Value
     * Term      intDocID  |  tf  |  Positions
     * apple     1            3      55, 98, 102
     */ 
	
	/**
	 * Write a batch of docs out into disk.
	 * @throws IOException
	 */
	private void writeABatchIntoDisk() throws IOException {
	    FileOutputStream fos = new FileOutputStream("data//" + dataType + ".index" + countOfBatches);
	    OutputStreamWriter osw = new OutputStreamWriter(fos);
	    bwForIndexFile = new BufferedWriter(osw);
	    
	    index.indexMap.forEach((k, v) -> { // Loop "term-postings lists" pairs in the indexMap
	        try {
	        	// 1. Write the current term(key) out 
	        	bwForIndexFile.write(k + " "); // k: The current term
	        	for (int i = 0; i < v.size(); i++) { // v: The current term's postings lists (A 2D array)
	        		// 2. Write the intDocId and tf(value) out 
	        		bwForIndexFile.write(v.get(i).get(0) + ":" + v.get(i).get(1) + ","); // v.get(i): Get every postings list(a doc)
	        																			 // get(0): intDocID | get(1): tf
	        		// 3. Write all positions(value) out 
	        		for(int j = 2; j < v.get(i).size(); j++) { // Write all positions                                                                		 
//	        			String s = (j == (v.get(i).size() - 1) ? ";" : ",");  
	        			String str = "";
	        			if(j < v.get(i).size()-1) { // If not reach the last position
	        				str = ","; // Split positions by ","
	        			} else { // If reach the last position 
	        				str = ";";  // End with ";"
	        			}
	        			
	        			bwForIndexFile.write(v.get(i).get(j) + str);
	        		} // End of inner for loop
	        	} // End of outer for loop
	        	
	        	// 4. Write all doc in the current batch and then start a newline
	        	/* ------- 每一个block写完后才换行，不懂！！！！！！！！！！！ ---------- */ 
	        	bwForIndexFile.write("\n"); 
	        } catch(IOException e) {
	        	e.printStackTrace();
	       	}
	    });
	    
	    countOfDocsPerBatch = 0; // Reset the count of docs processed per batch
	    countOfBatches++; // Increment 1 to track the next batch
	    index.clear();
	    bwForIndexFile.close();    
	}
	/* 
	 * What's inside the .index file: 
	 * Dictionary   Postings lists
	 * Term         IntegerDocID:TF,Positions;\n
	 * bee          1:3,55,67,88;\n
	 * can          1:4,66,90,106,170;\n
	 * ...
	 * 注意:TF后的标点可以改成";"
	 * 
	 * //Another way for loop a map
     * for(Entry<String, ArrayList<ArrayList<Integer>>> entry: ridxBlock.mapTerm.entrySet()) {
     *     String k = entry.getKey();
     *     ArrayList<ArrayList<Integer>> v = entry.getValue();
     * }
	 */
	
	/**
	 * Merge all batches of docs into a large index file
	 * @throws IOException
	 */
	private void mergeAllBatches() throws IOException {
		Map<String, String> mergedIndexMap = new HashMap<>(); // Key -> term | Value -> intDocID, tf, positions 
	    
		// 1. Read all batches of docs from separate index files
		for(int i = 0; i < countOfBatches; i++) { 
			 
			File separateIndexFile = new File("data//" + dataType + ".index" + i);
		    FileInputStream fis = new FileInputStream(separateIndexFile);
		    InputStreamReader isr = new InputStreamReader(fis);
		    BufferedReader br = new BufferedReader(isr);
		    
		    String str = "";
		    while((str = br.readLine()) != null) {
		    	 String[] sp = str.split("\\s"); // Split term(key) and postings list(value) up 
		    	 if(!mergedIndexMap.containsKey(sp[0])) { //sp[0]: key -> [term] 
		    		 mergedIndexMap.put(sp[0], sp[1]); //sp[1]: value -> [IntegerDocID]  
		    	 } else {
		    		  mergedIndexMap.put(sp[0], mergedIndexMap.get(sp[0]) + sp[1]); // String sp[0]: term | String sp[1]: postings list
		                                                                            // Append new position to the end of the postings list
		         }
		    }
		    
		    br.close(); // Close the br stream
		    separateIndexFile.delete(); // Delete the current separate index file after being written into the disk
		}
	
		// 2. Write all data in mergedIndexMap into the whole, large index file 
	    FileOutputStream fos = new FileOutputStream("data//" + dataType + ".index"); 																		
	    OutputStreamWriter osw = new OutputStreamWriter(fos);                                        
	    bwForIndexFile = new BufferedWriter(osw);
	    mergedIndexMap.forEach((k, v) -> {
	    	try {
	    		bwForIndexFile.write(k + " " + v + "\n");
	    	}catch (IOException e) {
	    		e.printStackTrace();
	    	}
	    });
	   
	    mergedIndexMap.clear(); // Clear the temporary mergedIndexMap for next batch writing
	    bwForIndexFile.close(); // Close the bwForIndexFile after all merges complete
	} 
	
	/**
	 * Create the dictionary file.
	 * @throws IOException
	 */
	private void createDictionary() throws IOException {
		// Read in the whole index file
		FileInputStream fis = new FileInputStream("data//" + dataType + ".index");
		InputStreamReader isr = new InputStreamReader(fis);
		BufferedReader br = new BufferedReader(isr);
		    
		// Write the dictinary out into disk
		//.dict文件是对.ridx的一个summary
		FileOutputStream fos = new FileOutputStream("data//" + dataType + ".dict");
		OutputStreamWriter osw = new OutputStreamWriter(fos);
		BufferedWriter bwForDictFile = new BufferedWriter(osw);
		  
		int countOfTerms = 0; //记录当前有多少行，可以算出最后总共有多少行
		String strLine = "";
		while((strLine = br.readLine()) != null) {
		    String[] term = strLine.split("\\s"); // Get the term part in the index file
		    bwForDictFile.write(term[0] + "," + countOfTerms++ + "\n"); // Assign each term an unique integer identifier		                                                              
		}
		 
		bwForDictFile.close();
		br.close();
	}
	/* 
	 * .dict structure:
	 * Terms   iLineNum
	 * apple   0
	 * bee     1
	 * candle  2
	 */
	
	/**
	 * Close the index writer. 
	 * All buffered content(if any) should be output.
	 * If you write your index into several files, you need to fuse them here.
	 * @throws IOException
	 */
	public void Close() throws IOException {
		// Write out into disk again before the program ends 
		writeABatchIntoDisk(); 
		
		// Merge all separate index files and create a dictionary file 
		mergeAllBatches(); 
		createDictionary(); 
		    
		// Clear the map of the basic structure of index
		index.clear();
		    
		// Close all BufferedWriter streams
		bwForDocIDFile.close();
	    bwForIndexFile.close();
	}
	
	/**
	 * Inner class: Index 
	 * The basic structure of the index. A mapping from the dictionary to its corresponding postings lists.
	 */
	private static final class Index {
	    
		private Map<String, ArrayList<ArrayList<Integer>>> indexMap; // A mapping from a term to its postings list
		                                                             // key: a term in the dictionary 
		                                                             // value: corresponding postings list
         
		/**
		 * Constructor: Initiate the whole index structure 
		 */
		public Index() {
		    indexMap = new HashMap<String, ArrayList<ArrayList<Integer>>>();
		}

		/**
		 * Update the index once a doc is added.
		 * @param doc
		 * @return
		 */
		private boolean update(Map<String, ArrayList<Integer>> doc) { 
		    if(doc != null) {
		        doc.forEach((k, v) -> {
		            if(!indexMap.containsKey(k)) {
		            	indexMap.put(k, new ArrayList<ArrayList<Integer>>(Arrays.asList(v)));
		            } else {
		            	indexMap.get(k).add(v);
		            }
		        });
		        
		        return true;
		    }
		    
		    return false;
	    }

		/**
		 * Clear the whole index structure
		 */
	    private void clear() {
	        if(!indexMap.isEmpty()) {
	    	    indexMap.clear();
	        }  
	    }
   	} // End of class Index
	/*
	 * indexMap
	 * key -> dictionary(terms)
	 * value -> postings lists(2D array)
	 * 
	 * Dictionary(Part I in index)      Postings lists(Part II in idex)                                  
	 * Term                             intDocID | tf | positions  
	 * apple                            1          2    55, 78            doc
	 *                                  7          4    23, 45, 78, 98    doc  
	 *                                  ...
	 *                                  100                               new doc                             
	 */
	
//	/**
//	 * Merge all batches of docs into a large index file
//	 * @throws IOException
//	 */
//	private void mergeAllBatches() throws IOException {
//		Map<String, String> mergedIndexMap = new HashMap<String, String>(); // Key -> term | Value -> intDocID, tf, positions 
//	    
//		FileOutputStream fos = new FileOutputStream("data//" + dataType + ".index"); 																		
//	    OutputStreamWriter osw = new OutputStreamWriter(fos);                                        
//	    bwForIndexFile = new BufferedWriter(osw);
//	   
//		for(int i = 0; i < countOfBatches; i++) { 
//			// 1. Read a batch of docs from the current separate index file 
//			File separateIndexFile = new File("data//" + dataType + ".index" + i);
//		    FileInputStream fis = new FileInputStream(separateIndexFile);
//		    InputStreamReader isr = new InputStreamReader(fis);
//		    BufferedReader br = new BufferedReader(isr);
//		    
//		    String str = "";
//		    while((str = br.readLine()) != null) {
//		    	 String[] sp = str.split("\\s"); // Split term(key) and postings list(value) up 
//		    	 if(!mergedIndexMap.containsKey(sp[0])) { //sp[0]: key -> [term] 
//		    		 mergedIndexMap.put(sp[0], sp[1]); //sp[1]: value -> [IntegerDocID]  
//		    	 } else {
//		    		  mergedIndexMap.put(sp[0], mergedIndexMap.get(sp[0]) + sp[1]); // String sp[0]: term | String sp[1]: postings list
//		                                                                            // Append new position to the end of the postings list
//		         }
//		    }
//		    
//		    br.close(); // Close the br stream
//	      
//		    // 2. Write the current batches of docs into the whole, large index file  
//		    mergedIndexMap.forEach((k, v) -> {
//		    	try {
//		    		bwForIndexFile.write(k + " " + v + "\n");
//		    	}catch (IOException e) {
//		    		e.printStackTrace();
//		    	}
//		    });
//		    
//		    // 3. Delete the current separate index file after being written into the disk 
//		    separateIndexFile.delete();
//		    
//		    // 4. Clear the temporary mergedIndexMap for next batch writing
//		    mergedIndexMap.clear();   
//		}
//	
//		// 5. Close the bwForIndexFile after all merges complete
//	    bwForIndexFile.close();
//	} 
}